import os
import numpy as np
import pandas as pd
import netCDF4 as nc
from datetime import datetime
from scipy.spatial import cKDTree
from read_CloudSat import reader

def load_fy_coordinates(coord_file_name):
    """加载风云卫星的地理坐标数据"""
    with nc.Dataset(coord_file_name, 'r') as coord_file:
        lat_fy = coord_file.variables['lat'][:, :].T
        lon_fy = coord_file.variables['lon'][:, :].T
        lat_fy[np.isnan(lat_fy)] = -90.
        lon_fy[np.isnan(lon_fy)] = 360.
    return lat_fy, lon_fy

def process_cloudsat_file(filepath):
    """处理CloudSat文件，提取地理数据、云属性和时间范围"""
    f = reader(filepath)
    lon_c, lat_c, elv = f.read_geo()
    height = f.read_sds('Height')
    cloud_mask = f.read_sds('CPR_Cloud_mask')
    time = f.read_time(datetime=True)
    cloudsat_start_time, cloudsat_end_time = time[[0, -1]]
    f.close()

    cloudsat_cth = []
    cloudsat_cbh = []
    cloud_types = []

    for i, row_data in enumerate(cloud_mask):
        if np.any(row_data >= 20):
            col_idx = np.argmax(row_data >= 20)
            last_idx = np.where(row_data >= 20)[0][-1]
            cloudsat_cth.append(height[i, col_idx] if col_idx < height.shape[1] else np.nan)
            cloudsat_cbh.append(height[i, last_idx] if last_idx < height.shape[1] else np.nan)
            cloud_types.append(1 if col_idx < last_idx and np.any(row_data[col_idx + 1:last_idx] < 20) else 0)
        else:
            cloudsat_cth.append(np.nan)
            cloudsat_cbh.append(np.nan)
            cloud_types.append(np.nan)

    return lon_c, lat_c, np.array(cloudsat_cth), np.array(cloudsat_cbh), cloudsat_start_time, cloudsat_end_time, np.array(cloud_types)

def match_cloudsat_with_fy(cloudsat_data, lat_fy, lon_fy, cth_fy_data, clm_fy_data, clt_fy_data, crf_fy_data, ctt_fy_data, ctp_fy_data, olr_fy_data, lpwlow_fy_data, lpwmid_fy_data, lpwhigh_fy_data, threshold=2):
    """匹配CloudSat数据与风云卫星数据"""
    lon_c, lat_c, cloudsat_cth, cloudsat_cbh, cloud_types = cloudsat_data
    threshold_deg = threshold / 111.32
    valid_mask = (cth_fy_data != 65535) & (clm_fy_data == 0)
    valid_lat_fy = lat_fy[valid_mask]
    valid_lon_fy = lon_fy[valid_mask]
    valid_cth_fy_data = cth_fy_data[valid_mask]
    valid_clm_fy_data = clm_fy_data[valid_mask]
    valid_clt_fy_data = clt_fy_data[valid_mask]
    valid_crf_fy_data = crf_fy_data[valid_mask]
    valid_ctt_fy_data = ctt_fy_data[valid_mask]
    valid_ctp_fy_data = ctp_fy_data[valid_mask]
    valid_olr_fy_data = olr_fy_data[valid_mask]
    valid_lpwlow_fy_data = lpwlow_fy_data[valid_mask]
    valid_lpwmid_fy_data = lpwmid_fy_data[valid_mask]
    valid_lpwhigh_fy_data = lpwhigh_fy_data[valid_mask]

    tree = cKDTree(np.c_[valid_lon_fy.ravel(), valid_lat_fy.ravel()])
    dist, idx = tree.query(np.c_[lon_c, lat_c], k=1, distance_upper_bound=threshold_deg)
    matched_indices = idx[dist != np.inf]
    matched_distances = dist[dist != np.inf]
    matched_cloudsat_points = np.arange(len(lat_c))[dist < threshold_deg]

    results = [{
        'cloudsat_idx': c_idx,
        'cloudsat_lat': lat_c[c_idx],
        'cloudsat_lon': lon_c[c_idx],
        'fy_lat': valid_lat_fy.flat[f_idx],
        'fy_lon': valid_lon_fy.flat[f_idx],
        'fy_cth': valid_cth_fy_data.flat[f_idx],
        'fy_clm': valid_clm_fy_data.flat[f_idx],
        'fy_clt': valid_clt_fy_data.flat[f_idx],
        'fy_crf': valid_crf_fy_data.flat[f_idx],
        'fy_ctt': valid_ctt_fy_data.flat[f_idx],
        'fy_ctp': valid_ctp_fy_data.flat[f_idx],
        'fy_olr': valid_olr_fy_data.flat[f_idx],
        'fy_lpwlow': valid_lpwlow_fy_data.flat[f_idx],
        'fy_lpwmid': valid_lpwmid_fy_data.flat[f_idx],
        'fy_lpwhigh': valid_lpwhigh_fy_data.flat[f_idx],
        'cloudsat_cth': cloudsat_cth[c_idx],
        'cloudsat_cbh': cloudsat_cbh[c_idx],
        'cloudsat_tpye': cloud_types[c_idx],
        'distance': d
    } for c_idx, f_idx, d in zip(matched_cloudsat_points, matched_indices, matched_distances)]

    closest_matches = {}
    for result in results:
        fy_key = (result['fy_lat'], result['fy_lon'])
        if fy_key not in closest_matches or result['distance'] < closest_matches[fy_key]['distance']:
            closest_matches[fy_key] = result

    return [result for result in closest_matches.values() if result['cloudsat_cth'] != result['cloudsat_cbh'] and -1000 <= result['fy_cth'] - result['cloudsat_cth'] <= 1000]

def main():
    cloudsat_filepath = '/mnt/space2/liudd/cloudsatdate/2019048060356_68216_CS_2B-GEOPROF_GRANULE_P1_R05_E08_F03.hdf'
    fy_filepaths = {
        'cth': '/home/liudd/data/fy4a_cth/FY4A-_AGRI--_N_DISK_1047E_L2-_CTH-_MULT_NOM_20190217070000_20190217071459_4000M_V0001.NC',
        'clm': '/home/liudd/data/fy4a_clm/FY4A-_AGRI--_N_DISK_1047E_L2-_CLM-_MULT_NOM_20190217070000_20190217071459_4000M_V0001.NC',
        'clt': '/home/liudd/data/fy4a_clt/FY4A-_AGRI--_N_DISK_1047E_L2-_CLT-_MULT_NOM_20190217070000_20190217071459_4000M_V0001.NC',
        'crf': '/home/liudd/data/fy4a_crf/FY4A-_AGRI--_N_DISK_1047E_L2-_CFR-_MULT_NOM_20190217070000_20190217071459_4000M_V0001.NC',
        'ctt': '/home/liudd/data/fy4a_ctt/FY4A-_AGRI--_N_DISK_1047E_L2-_CTT-_MULT_NOM_20190217070000_20190217071459_4000M_V0001.NC',
        'ctp': '/home/liudd/data/fy4a_ctp/FY4A-_AGRI--_N_DISK_1047E_L2-_CTP-_MULT_NOM_20190217070000_20190217071459_4000M_V0001.NC',
        'olr': '/home/liudd/data/fy4a_olr/FY4A-_AGRI--_N_DISK_1047E_L2-_OLR-_MULT_NOM_20190217070000_20190217071459_4000M_V0001.NC',
        'lpw': '/home/liudd/data/fy4a_lpw/FY4A-_AGRI--_N_DISK_1047E_L2-_LPW-_MULT_NOM_20190217070000_20190217071459_4000M_V0001.NC',
    }
    coord_file_name = 'FY4A_coordinates.nc'

    lat_fy, lon_fy = load_fy_coordinates(coord_file_name)
    lon_c, lat_c, cloudsat_cth, cloudsat_cbh, cloudsat_start_time, cloudsat_end_time, cloudsat_types = process_cloudsat_file(cloudsat_filepath)

    with nc.Dataset(fy_filepaths['cth'], 'r') as cth_file:
        cth_data = cth_file.variables['CTH'][:]
    with nc.Dataset(fy_filepaths['clm'], 'r') as clm_file:
        clm_data = clm_file.variables['CLM'][:]
    with nc.Dataset(fy_filepaths['clt'], 'r') as clt_file:
        clt_data = clt_file.variables['CLT'][:]
    with nc.Dataset(fy_filepaths['crf'], 'r') as crf_file:
        crf_data = crf_file.variables['CFR'][:]
    with nc.Dataset(fy_filepaths['ctt'], 'r') as ctt_file:
        ctt_data = ctt_file.variables['CTT'][:]
    with nc.Dataset(fy_filepaths['ctp'], 'r') as ctp_file:
        ctp_data = ctp_file.variables['CTP'][:]
    with nc.Dataset(fy_filepaths['olr'], 'r') as olr_file:
        olr_data = olr_file.variables['OLR'][:]
    with nc.Dataset(fy_filepaths['lpw'], 'r') as lpw_file:
        lpwlow_data = lpw_file.variables['LPW_LOW'][:]
        lpwmid_data = lpw_file.variables['LPW_MID'][:]
        lpwhigh_data = lpw_file.variables['LPW_HIGH'][:]

    match_results = match_cloudsat_with_fy(
        (lon_c, lat_c, cloudsat_cth, cloudsat_cbh, cloudsat_types),
        lat_fy, lon_fy, cth_data, clm_data, clt_data, crf_data, ctt_data, ctp_data, olr_data, lpwlow_data, lpwmid_data, lpwhigh_data)

    results_df = pd.DataFrame(match_results)
    results_df.to_csv('cloudsat_fy4a_matches.csv', index=False)

if __name__ == '__main__':
    main()
